﻿// Written By: Doug Finke
// Blog: http://dougfinke.com/blog

using System;
using System.IO;
using System.Runtime.InteropServices;
using System.Text;
using Microsoft.VisualStudio.TextTemplating.VSHost;
using Microsoft.Win32;

namespace PowerShellCustomTool
{

    [Guid("cb65c781-2f19-47c1-9fc6-b035a32c0440")]
    [ComVisible(true)]
    public class PowerShellCustomTool : BaseCodeGeneratorWithSite
    {
        string _fileExtension = ".cs";

        protected override byte[] GenerateCode(string inputFileName, string inputFileContent)
        {
            if (inputFileContent.Length > 0) 
            {
                PowerShellUtil.ThePowerShellEngine().Runspace = PowerShellUtil.TheRunspace();

                var thePSCGProfile = PowerShellUtil.ThePSCGProfile();
                PowerShellUtil.TheRunspace().SessionStateProxy.SetVariable("PSCGProfile", thePSCGProfile);
                PowerShellUtil.TheRunspace().SessionStateProxy.SetVariable("DTE", this.Dte);

                // Source the Profile
                if (File.Exists(thePSCGProfile))
                {
                    File.ReadAllText(thePSCGProfile).InvokePowerShell(out _fileExtension);
                }

                // Run the PowerShell code in the file passed from Visual Studio
                //var result = (inputFileContent.Trim() + " | Out-String -Width 133").InvokePowerShell(out _fileExtension);
                var result = (inputFileContent.Trim()).InvokePowerShell(out _fileExtension);
                
                // Dispose/Close PowerShell and the Runspace
                PowerShellUtil.ReleasePowerShell();

                return Encoding.UTF8.GetBytes(result);
            }

            return null;
        }

        public override string GetDefaultExtension()
        {
            return _fileExtension;
        }

        #region COM Registration
        //    Windows Registry Editor Version 5.00
        // [HKEY_LOCAL_MACHINE\SOFTWARE\Microsoft\VisualStudio\10.0\Generators\{FAE04EC1-301F-11D3-BF4B-00C04F79EFBC}\PSCodeGen]
        // @="PowerShell Class Generator"
        // "CLSID"="{cb65c781-2f19-47c1-9fc6-b035a32c0440}"
        // "GeneratesDesignTimeSource"=dword:00000001

        private static Guid CustomToolGuid = new Guid("{CB65C781-2F19-47C1-9FC6-B035A32C0440}");
        private static Guid CSharpCategory = new Guid("{FAE04EC1-301F-11D3-BF4B-00C04F79EFBC}");
        private static Guid VBCategory = new Guid("{164B10B9-B200-11D0-8C61-00A0C91E29D5}");
        
        private const string CustomToolName = "PSCodeGenerator";
        private const string CustomToolDescription = "PowerShell Class Generator";
        private const string KeyFormat = @"SOFTWARE\Microsoft\VisualStudio\{0}\Generators\{1}\{2}";

        protected static void Register(Version vsVersion, Guid categoryGuid)
        {
            Console.WriteLine(CustomToolGuid.ToString("B"));
            Console.WriteLine(String.Format(KeyFormat, vsVersion, categoryGuid.ToString("B"), CustomToolName));

            using (RegistryKey key = Registry.LocalMachine.CreateSubKey(String.Format(KeyFormat, vsVersion, categoryGuid.ToString("B"), CustomToolName)))
            {
                key.SetValue("", CustomToolDescription);
                key.SetValue("CLSID", CustomToolGuid.ToString("B"));
                key.SetValue("GeneratesDesignTimeSource", 1);
            }
        }

        protected static void UnRegister(Version vsVersion, Guid categoryGuid)
        {
            Console.WriteLine(String.Format(KeyFormat, vsVersion, categoryGuid.ToString("B"), CustomToolName));
            Registry.LocalMachine.DeleteSubKey(
                String.Format(KeyFormat, vsVersion, categoryGuid.ToString("B"), CustomToolName), false);
        }

        [ComRegisterFunction]
        public static void RegisterClass(Type t)
        {
            // Register for both VS.NET 2010 (C#)
            Register(new Version(10, 0), CSharpCategory);

            // Register for both VS.NET 2010 (VB)
            //Register(new Version(10, 0), VBCategory);
        }

        [ComUnregisterFunction]
        public static void UnRegisterClass(Type t)
        {
            // Register for both VS.NET 2010 (C#)
            UnRegister(new Version(10, 0), CSharpCategory);

            // Register for both VS.NET 2010 (VB)
            //UnRegister(new Version(10, 0), VBCategory);
        } 
        #endregion
    }
}